import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from LatticeInfo import *
import pickle

class RubikLattice:
    '''spin: 1表示顶面，6表示底面，2,3,4,5表示侧面，其中2位于左侧，3位于主视侧——主视：x轴正侧看向负方向'''
    def __init__(self,pos):
        if(len(pos)==3):
            self.initialPos=pos
            self.spin=[1,2,3,4,5,6] #表示自身的旋转状态
        else:
            print("position length error!")
            return
        
    def getInitialPos(self):
        return self.initialPos

class Rubik:
    def __init__(self):
        tempList=[]
        for i in range(3**3):
            tempList.append(RubikLattice([-1,-1,-1]))
        self.arrangeState=np.array(tempList).reshape((3,3,3))
        #self.arrangeState=np.zeros((3,3,3))
        for i in range(3):
            for j in range(3):
                for k in range(3):
                    self.arrangeState[i,j,k]=RubikLattice([i+1,j+1,k+1])
                    
    def rotate(self,plane,time=1):
        def perm(array,moveIndex,time): #旋转使得spin移位，time正数表示右移位，反之左移位
            if(len(moveIndex)!=4):
                print("长度错误！")
                return
            else:
                p1=moveIndex[0]-1
                p2=moveIndex[1]-1
                p3=moveIndex[2]-1
                p4=moveIndex[3]-1
                if(time>0):
                    for i in range(time):
                        temp=array[p2]
                        array[p2]=array[p1]
                        array[p1]=array[p4]
                        array[p4]=array[p3]
                        array[p3]=temp
                else:
                    for i in range(-time):
                        temp=array[p2]
                        array[p2]=array[p3]
                        array[p3]=array[p4]
                        array[p4]=array[p1]
                        array[p1]=temp                   
            
        '''默认：顺时针（从坐标轴正半轴处往负轴方向来看）转一次'''
        if(len(plane)!=3):
            print("长度错误！")
            return
        if(type(time)!=int):
            print("旋转次数必须是整数！")
            return
        #get slice
        if(plane[0]!=0):
            rubikSlice=self.arrangeState[plane[0]-1,:,:]
        elif(plane[1]!=0):
            rubikSlice=self.arrangeState[:,plane[1]-1,:]
        else:
            rubikSlice=self.arrangeState[:,:,plane[2]-1]
        #rotate slice
        N=abs(time)%4 #normalize
        if(plane[1]!=0):
            if(time>=0): #顺时针
                rubikSlice=np.rot90(rubikSlice,N)
            else: #逆时针
                rubikSlice=np.rot90(rubikSlice,-N)
        else:
            if(time>=0): #顺时针
                rubikSlice=np.rot90(rubikSlice,-N)
            else: #逆时针
                rubikSlice=np.rot90(rubikSlice,N)
        #传回
        if(plane[0]!=0):
            self.arrangeState[plane[0]-1,:,:]=rubikSlice
        elif(plane[1]!=0):
            self.arrangeState[:,plane[1]-1,:]=rubikSlice
        else:
            self.arrangeState[:,:,plane[2]-1]=rubikSlice
        #更新各方块的spin
        if(plane[0]!=0): #转的x轴
            for subLayerRubikSlice in rubikSlice:
                for cube in subLayerRubikSlice:
                    perm(cube.spin,[1,4,6,2],time)
        elif(plane[1]!=0): #转的y轴
            for subLayerRubikSlice in rubikSlice:
                for cube in subLayerRubikSlice:
                    perm(cube.spin,[1,5,6,3],time)
        else:
            for subLayerRubikSlice in rubikSlice:
                for cube in subLayerRubikSlice:
                    perm(cube.spin,[5,4,3,2],time)
    
    def showState(self):
        print("z\n•→y\n↓\nx")
        print("z=1:")
        for i in range(3):
            for j in range(3):
                print(self.arrangeState[i,j,0].getInitialPos(),end=" ")
            print("")
        print("z=2:")
        for i in range(3):
            for j in range(3):
                print(self.arrangeState[i,j,1].getInitialPos(),end=" ")
            print("")
        print("z=3:")
        for i in range(3):
            for j in range(3):
                print(self.arrangeState[i,j,2].getInitialPos(),end=" ")
            print("")
            
    def showImg(self):
        def plotPlane(pos,planeFlag,isBlock,fig,ax,Color=None): #画出平面，如果不是画方格，就画底版,planeFlag：决定平面所在的面
            #作图原点
            x0=pos[0]
            y0=pos[1]
            z0=pos[2]
            if(not isBlock): #底版模式
                L=0.5
                x=np.array([0,L])
                y=np.array([0,L])
                z=np.array([0,L])
                if(planeFlag[0]!=0): #x平面
                    y,z=np.meshgrid(y,z)
                    x=x0+y*0
                elif(planeFlag[1]!=0): #y平面
                    z,x=np.meshgrid(z,x)
                    y=y0+z*0
                else: #z平面
                    x,y=np.meshgrid(x,y)
                    z=z0+x*0
                test=ax.plot_surface(x,y,z,rstride=1,cstride=1,color=[0,0,0])
                test.set_zorder(1)
                ax.add_collection3d(test)
            else: #方格模式
                L=1
                dL=0.01 #边框
                x=np.array([x0+dL,x0+L-dL])
                y=np.array([y0+dL,y0+L-dL])
                z=np.array([z0+dL,z0+L-dL])
                if(planeFlag[0]!=0): #x平面
                    y,z=np.meshgrid(y,z)
                    x=x0+y*0
                elif(planeFlag[1]!=0): #y平面
                    z,x=np.meshgrid(z,x)
                    y=y0+z*0
                else: #z平面
                    x,y=np.meshgrid(x,y)
                    z=z0+x*0
                test=ax.plot_surface(x,y,z,rstride=1,cstride=1,color=Color)
                test.set_zorder(10)
                ax.add_collection3d(test)
            
        '''画图'''
        fig,ax=plt.subplots(1,1,figsize=(7,6))
        ax=Axes3D(fig)
        #画出底板
        #plotPlane([0,0,0],[1,0,0],False,fig,ax)
        #plotPlane([0,0,0],[0,1,0],False,fig,ax)
        #plotPlane([0,0,0],[0,0,1],False,fig,ax)
        #plotPlane([0,0,3],[0,0,1],False,fig,ax)
        #plotPlane([3,0,0],[1,0,0],False,fig,ax)
        #plotPlane([0,3,0],[0,1,0],False,fig,ax)
        #画出六个面
        #1:z=3
        slice1=self.arrangeState[:,:,3-1]
        for i in range(3):
            for j in range(3):
                block=slice1[i,j]
                blockCurrentFace=block.spin[1-1]
                plotPlane([i,j,3],[0,0,1],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #2:y=1
        slice2=self.arrangeState[:,1-1,:]
        for i in range(3):
            for j in range(3):
                block=slice2[i,j]
                blockCurrentFace=block.spin[2-1]
                plotPlane([i,0,j],[0,1,0],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #3:x=3
        slice3=self.arrangeState[3-1,:,:]
        for i in range(3):
            for j in range(3):
                block=slice3[i,j]
                blockCurrentFace=block.spin[3-1]
                plotPlane([3,i,j],[1,0,0],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #4:y=3
        slice4=self.arrangeState[:,3-1,:]
        for i in range(3):
            for j in range(3):
                block=slice4[i,j]
                blockCurrentFace=block.spin[4-1]
                plotPlane([i,3,j],[0,1,0],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #5:x=1
        slice5=self.arrangeState[1-1,:,:]
        for i in range(3):
            for j in range(3):
                block=slice5[i,j]
                blockCurrentFace=block.spin[5-1]
                plotPlane([0,i,j],[1,0,0],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #6:z=1
        slice5=self.arrangeState[:,:,1-1]
        for i in range(3):
            for j in range(3):
                block=slice5[i,j]
                blockCurrentFace=block.spin[6-1]
                plotPlane([i,j,0],[0,0,1],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #画出网格
        ax.plot([0,0,3,3,0],[0,3,3,0,0],[3,3,3,3,3],'k-')
        ax.plot([0,0,3,3,0],[0,3,3,0,0],[0,0,0,0,0],'k-')
        ax.plot([0,0,3,3,0],[0,3,3,0,0],[1,1,1,1,1],'k-')
        ax.plot([0,0,3,3,0],[0,3,3,0,0],[2,2,2,2,2],'k-')
        ax.plot([0,0,3,3,0],[3,3,3,3,3],[0,3,3,0,0],'k-')
        ax.plot([0,0,3,3,0],[0,0,0,0,0],[0,3,3,0,0],'k-')
        ax.plot([0,0,3,3,0],[1,1,1,1,1],[0,3,3,0,0],'k-')
        ax.plot([0,0,3,3,0],[2,2,2,2,2],[0,3,3,0,0],'k-')
        ax.plot([1,1,1,1,1],[0,0,3,3,0],[0,3,3,0,0],'k-')
        ax.plot([2,2,2,2,2],[0,0,3,3,0],[0,3,3,0,0],'k-')
        ##########
        ax.axis('off')
        ax.view_init(45,45)
        fig.show()
        
class RubikInteraction:
    '''图形界面交互式'''
    def __init__(self):
        tempList=[]
        for i in range(3**3):
            tempList.append(RubikLattice([-1,-1,-1]))
        self.arrangeState=np.array(tempList).reshape((3,3,3))
        #self.arrangeState=np.zeros((3,3,3))
        for i in range(3):
            for j in range(3):
                for k in range(3):
                    self.arrangeState[i,j,k]=RubikLattice([i+1,j+1,k+1])
                    
    def rotate(self,plane,time=1):
        def perm(array,moveIndex,time): #旋转使得spin移位，time正数表示右移位，反之左移位
            if(len(moveIndex)!=4):
                print("长度错误！")
                return
            else:
                p1=moveIndex[0]-1
                p2=moveIndex[1]-1
                p3=moveIndex[2]-1
                p4=moveIndex[3]-1
                if(time>0):
                    for i in range(time):
                        temp=array[p2]
                        array[p2]=array[p1]
                        array[p1]=array[p4]
                        array[p4]=array[p3]
                        array[p3]=temp
                else:
                    for i in range(-time):
                        temp=array[p2]
                        array[p2]=array[p3]
                        array[p3]=array[p4]
                        array[p4]=array[p1]
                        array[p1]=temp                   
            
        '''默认：顺时针（从坐标轴正半轴处往负轴方向来看）转一次'''
        if(len(plane)!=3):
            print("长度错误！")
            return
        if(type(time)!=int):
            print("旋转次数必须是整数！")
            return
        #get slice
        if(plane[0]!=0):
            rubikSlice=self.arrangeState[plane[0]-1,:,:]
        elif(plane[1]!=0):
            rubikSlice=self.arrangeState[:,plane[1]-1,:]
        else:
            rubikSlice=self.arrangeState[:,:,plane[2]-1]
        #rotate slice
        N=abs(time)%4 #normalize
        if(plane[1]!=0):
            if(time>=0): #顺时针
                rubikSlice=np.rot90(rubikSlice,N)
            else: #逆时针
                rubikSlice=np.rot90(rubikSlice,-N)
        else:
            if(time>=0): #顺时针
                rubikSlice=np.rot90(rubikSlice,-N)
            else: #逆时针
                rubikSlice=np.rot90(rubikSlice,N)
        #传回
        if(plane[0]!=0):
            self.arrangeState[plane[0]-1,:,:]=rubikSlice
        elif(plane[1]!=0):
            self.arrangeState[:,plane[1]-1,:]=rubikSlice
        else:
            self.arrangeState[:,:,plane[2]-1]=rubikSlice
        #更新各方块的spin
        if(plane[0]!=0): #转的x轴
            for subLayerRubikSlice in rubikSlice:
                for cube in subLayerRubikSlice:
                    perm(cube.spin,[1,4,6,2],time)
        elif(plane[1]!=0): #转的y轴
            for subLayerRubikSlice in rubikSlice:
                for cube in subLayerRubikSlice:
                    perm(cube.spin,[1,5,6,3],time)
        else:
            for subLayerRubikSlice in rubikSlice:
                for cube in subLayerRubikSlice:
                    perm(cube.spin,[5,4,3,2],time)
    
    def showState(self):
        print("z\n•→y\n↓\nx")
        print("z=1:")
        for i in range(3):
            for j in range(3):
                print(self.arrangeState[i,j,0].getInitialPos(),end=" ")
            print("")
        print("z=2:")
        for i in range(3):
            for j in range(3):
                print(self.arrangeState[i,j,1].getInitialPos(),end=" ")
            print("")
        print("z=3:")
        for i in range(3):
            for j in range(3):
                print(self.arrangeState[i,j,2].getInitialPos(),end=" ")
            print("")
            
    def getImg(self):
        def plotPlane(pos,planeFlag,isBlock,fig,ax,Color=None): #画出平面，如果不是画方格，就画底版,planeFlag：决定平面所在的面
            #作图原点
            x0=pos[0]
            y0=pos[1]
            z0=pos[2]
            if(not isBlock): #底版模式
                L=0.5
                x=np.array([0,L])
                y=np.array([0,L])
                z=np.array([0,L])
                if(planeFlag[0]!=0): #x平面
                    y,z=np.meshgrid(y,z)
                    x=x0+y*0
                elif(planeFlag[1]!=0): #y平面
                    z,x=np.meshgrid(z,x)
                    y=y0+z*0
                else: #z平面
                    x,y=np.meshgrid(x,y)
                    z=z0+x*0
                test=ax.plot_surface(x,y,z,rstride=1,cstride=1,color=[0,0,0])
                test.set_zorder(1)
                ax.add_collection3d(test)
            else: #方格模式
                L=1
                dL=0.01 #边框
                x=np.array([x0+dL,x0+L-dL])
                y=np.array([y0+dL,y0+L-dL])
                z=np.array([z0+dL,z0+L-dL])
                if(planeFlag[0]!=0): #x平面
                    y,z=np.meshgrid(y,z)
                    x=x0+y*0
                elif(planeFlag[1]!=0): #y平面
                    z,x=np.meshgrid(z,x)
                    y=y0+z*0
                else: #z平面
                    x,y=np.meshgrid(x,y)
                    z=z0+x*0
                test=ax.plot_surface(x,y,z,rstride=1,cstride=1,color=Color)
                test.set_zorder(10)
                ax.add_collection3d(test)
            
        '''画图'''
        fig,ax=plt.subplots(1,1,figsize=(7,6))
        ax=Axes3D(fig)
        #画出底板
        #plotPlane([0,0,0],[1,0,0],False,fig,ax)
        #plotPlane([0,0,0],[0,1,0],False,fig,ax)
        #plotPlane([0,0,0],[0,0,1],False,fig,ax)
        #plotPlane([0,0,3],[0,0,1],False,fig,ax)
        #plotPlane([3,0,0],[1,0,0],False,fig,ax)
        #plotPlane([0,3,0],[0,1,0],False,fig,ax)
        #画出六个面
        #1:z=3
        slice1=self.arrangeState[:,:,3-1]
        for i in range(3):
            for j in range(3):
                block=slice1[i,j]
                blockCurrentFace=block.spin[1-1]
                plotPlane([i,j,3],[0,0,1],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #2:y=1
        slice2=self.arrangeState[:,1-1,:]
        for i in range(3):
            for j in range(3):
                block=slice2[i,j]
                blockCurrentFace=block.spin[2-1]
                plotPlane([i,0,j],[0,1,0],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #3:x=3
        slice3=self.arrangeState[3-1,:,:]
        for i in range(3):
            for j in range(3):
                block=slice3[i,j]
                blockCurrentFace=block.spin[3-1]
                plotPlane([3,i,j],[1,0,0],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #4:y=3
        slice4=self.arrangeState[:,3-1,:]
        for i in range(3):
            for j in range(3):
                block=slice4[i,j]
                blockCurrentFace=block.spin[4-1]
                plotPlane([i,3,j],[0,1,0],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #5:x=1
        slice5=self.arrangeState[1-1,:,:]
        for i in range(3):
            for j in range(3):
                block=slice5[i,j]
                blockCurrentFace=block.spin[5-1]
                plotPlane([0,i,j],[1,0,0],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #6:z=1
        slice5=self.arrangeState[:,:,1-1]
        for i in range(3):
            for j in range(3):
                block=slice5[i,j]
                blockCurrentFace=block.spin[6-1]
                plotPlane([i,j,0],[0,0,1],True,fig,ax,Color=LatticeInfo.blockColor[str(block.getInitialPos())][blockCurrentFace-1])
        #画出网格
        ax.plot([0,0,3,3,0],[0,3,3,0,0],[3,3,3,3,3],'k-')
        ax.plot([0,0,3,3,0],[0,3,3,0,0],[0,0,0,0,0],'k-')
        ax.plot([0,0,3,3,0],[0,3,3,0,0],[1,1,1,1,1],'k-')
        ax.plot([0,0,3,3,0],[0,3,3,0,0],[2,2,2,2,2],'k-')
        ax.plot([0,0,3,3,0],[3,3,3,3,3],[0,3,3,0,0],'k-')
        ax.plot([0,0,3,3,0],[0,0,0,0,0],[0,3,3,0,0],'k-')
        ax.plot([0,0,3,3,0],[1,1,1,1,1],[0,3,3,0,0],'k-')
        ax.plot([0,0,3,3,0],[2,2,2,2,2],[0,3,3,0,0],'k-')
        ax.plot([1,1,1,1,1],[0,0,3,3,0],[0,3,3,0,0],'k-')
        ax.plot([2,2,2,2,2],[0,0,3,3,0],[0,3,3,0,0],'k-')
        ##########
        ax.axis('off')
        ax.view_init(45,45)
        return fig,ax
    
    def writeAsFile(self,name="default.rubik"):
        with open(name,"wb") as f:
            pickle.dump(self.arrangeState,f)
            
    def loadFile(self,name="default.rubik"):
        with open(name,"rb") as f:
            self.arrangeState=pickle.load(f)
            
    def getState(self):
        return self.arrangeState
    
    def setState(self,st):
        self.arrangeState=st